from datetime import datetime

from sqlalchemy import Column, Integer, String, DateTime, Text, Enum, ForeignKey, DECIMAL
from sqlalchemy.orm import relationship
from app.extensions import db


class TreatmentRecord(db.Model):
    __tablename__ = 'treatment_records'
    
    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, ForeignKey('users.id'), nullable=False)
    device_id = Column(String(50))
    treatment_mode = Column(Enum('tDCS+', 'tDCS-', 'tACS'), nullable=False)
    current_intensity = Column(DECIMAL(3, 2))  # 0-2.0mA
    treatment_duration = Column(Integer)  # 分钟
    frequency = Column(DECIMAL(5, 1))  # Hz (仅tACS模式)
    start_time = Column(DateTime)
    end_time = Column(DateTime)
    status = Column(Enum('completed', 'interrupted', 'cancelled'), default='completed')
    effectiveness_score = Column(Integer)  # 1-10评分
    notes = Column(Text)
    created_at = Column(DateTime, default=datetime.now)
    
    def to_dict(self):
        return {
            'id': self.id,
            'user_id': self.user_id,
            'device_id': self.device_id,
            'treatment_mode': self.treatment_mode,
            'current_intensity': float(self.current_intensity) if self.current_intensity else None,
            'treatment_duration': self.treatment_duration,
            'frequency': float(self.frequency) if self.frequency else None,
            'start_time': self.start_time.isoformat() if self.start_time else None,
            'end_time': self.end_time.isoformat() if self.end_time else None,
            'status': self.status,
            'effectiveness_score': self.effectiveness_score,
            'notes': self.notes,
            'created_at': self.created_at.isoformat() if self.created_at else None,
            'actual_duration': self._calculate_actual_duration()
        }
    
    def _calculate_actual_duration(self):
        """计算实际治疗时长（分钟，保留两位小数）"""
        if self.start_time and self.end_time:
            delta = self.end_time - self.start_time
            return round(delta.total_seconds() / 60, 2)
        return None


class TreatmentTemplate(db.Model):
    __tablename__ = 'treatment_templates'
    
    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, ForeignKey('users.id'), nullable=False)
    template_name = Column(String(100), nullable=False)
    template_type = Column(Enum('hospital', 'school', 'enterprise', 'research'), default='hospital')
    treatment_mode = Column(Enum('tDCS+', 'tDCS-', 'tACS'), nullable=False)
    current_intensity = Column(DECIMAL(3, 2))  # 改为支持0.01mA精度
    treatment_duration = Column(Integer)
    frequency = Column(DECIMAL(5, 1))
    is_default = Column(db.Boolean, default=False)
    created_at = Column(DateTime, default=datetime.now)
    
    def to_dict(self):
        return {
            'id': self.id,
            'user_id': self.user_id,
            'template_name': self.template_name,
            'template_type': self.template_type,
            'treatment_mode': self.treatment_mode,
            'current_intensity': float(self.current_intensity) if self.current_intensity else None,
            'treatment_duration': self.treatment_duration,
            'frequency': float(self.frequency) if self.frequency else None,
            'is_default': self.is_default,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }
